package demo;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.net.URLEncoder;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.stream.FactoryConfigurationError;
import javax.xml.stream.XMLStreamException;

import org.joda.time.DateTime;
import org.opensaml.Configuration;
import org.opensaml.DefaultBootstrap;
import org.opensaml.common.SAMLVersion;
import org.opensaml.saml2.core.AuthnContextClassRef;
import org.opensaml.saml2.core.AuthnContextComparisonTypeEnumeration;
import org.opensaml.saml2.core.AuthnRequest;
import org.opensaml.saml2.core.Issuer;
import org.opensaml.saml2.core.NameIDPolicy;
import org.opensaml.saml2.core.RequestedAuthnContext;
import org.opensaml.saml2.core.impl.AuthnContextClassRefBuilder;
import org.opensaml.saml2.core.impl.AuthnRequestBuilder;
import org.opensaml.saml2.core.impl.IssuerBuilder;
import org.opensaml.saml2.core.impl.NameIDPolicyBuilder;
import org.opensaml.saml2.core.impl.RequestedAuthnContextBuilder;
import org.opensaml.xml.ConfigurationException;
import org.opensaml.xml.XMLObject;
import org.opensaml.xml.io.Marshaller;
import org.opensaml.xml.io.MarshallingException;
import org.opensaml.xml.io.Unmarshaller;
import org.opensaml.xml.io.UnmarshallerFactory;
import org.opensaml.xml.io.UnmarshallingException;
import org.opensaml.xml.util.Base64;
import org.opensaml.xml.util.XMLHelper;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public class SamlConsumer {

	private String identityProviderUrl = null;
	private String resourceProviderUrl = null;
	private String consumerUrl = null;
	private String randomId = Integer.toHexString(new Double(Math.random())
			.intValue());
	private String relayState = null;

	public String buildAuthnRequestMessage(String resourceUrl) {

		AuthnRequest authnRequest = null;
		String encodedAuthRequest = null;
		resourceProviderUrl = resourceUrl;
		consumerUrl = resourceUrl;
		String redirectUrl = null;

		try {

			DefaultBootstrap.bootstrap();

			identityProviderUrl = new ConfigFileReader().getIssuerUrl().trim();

			authnRequest = this.buildAuthnRequestObject();// SAML AuthRequest

			encodedAuthRequest = encodeAuthnRequest(authnRequest);

		} catch (ConfigurationException e) {
			e.printStackTrace();

		} catch (XMLStreamException ex) {
			ex.printStackTrace();

		} catch (FactoryConfigurationError exc) {
			exc.printStackTrace();

		} catch (MarshallingException e) {
			e.printStackTrace();

		} catch (IOException e) {
			e.printStackTrace();

		}

		redirectUrl = identityProviderUrl + "?SAMLRequest="
				+ encodedAuthRequest + "&RelayState=" + relayState;

		return redirectUrl;
	}

	private AuthnRequest buildAuthnRequestObject() {

		IssuerBuilder issuerBuilder = null;
		Issuer issuer = null;
		NameIDPolicyBuilder nameIdPolicyBuilder = null;
		NameIDPolicy nameIdPolicy = null;
		AuthnContextClassRefBuilder authnContextClassRefBuilder = null;
		AuthnContextClassRef authnContextClassRef = null;
		RequestedAuthnContextBuilder requestedAuthnContextBuilder = null;
		RequestedAuthnContext requestedAuthnContext = null;
		DateTime issueInstant = null;
		AuthnRequestBuilder authRequestBuilder = null;
		AuthnRequest authRequest = null;

		// Issuer object
		issuerBuilder = new IssuerBuilder();
		issuer = issuerBuilder.buildObject(
				"urn:oasis:names:tc:SAML:2.0:assertion", "Issuer", "samlp");
		issuer.setValue(resourceProviderUrl);

		// NameIDPolicy
		nameIdPolicyBuilder = new NameIDPolicyBuilder();
		nameIdPolicy = nameIdPolicyBuilder.buildObject();
		nameIdPolicy
				.setFormat("urn:oasis:names:tc:SAML:2.0:nameid-format:persistent");
		nameIdPolicy.setSPNameQualifier("Isser");
		nameIdPolicy.setAllowCreate(new Boolean(true));

		// AuthnContextClass
		authnContextClassRefBuilder = new AuthnContextClassRefBuilder();
		authnContextClassRef = authnContextClassRefBuilder.buildObject(
				"urn:oasis:names:tc:SAML:2.0:assertion",
				"AuthnContextClassRef", "saml");
		authnContextClassRef
				.setAuthnContextClassRef("urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport");

		// AuthnContex
		requestedAuthnContextBuilder = new RequestedAuthnContextBuilder();
		requestedAuthnContext = requestedAuthnContextBuilder.buildObject();
		requestedAuthnContext
				.setComparison(AuthnContextComparisonTypeEnumeration.EXACT);
		requestedAuthnContext.getAuthnContextClassRefs().add(
				authnContextClassRef);

		// Creation of AuthRequestObject
		issueInstant = new DateTime();
		authRequestBuilder = new AuthnRequestBuilder();
		authRequest = authRequestBuilder
				.buildObject("urn:oasis:names:tc:SAML:2.0:protocol",
						"AuthnRequest", "samlp");
		authRequest.setForceAuthn(new Boolean(false));
		authRequest.setIsPassive(new Boolean(false));
		authRequest.setIssueInstant(issueInstant);
		authRequest
				.setProtocolBinding("urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST");
		authRequest.setAssertionConsumerServiceURL(consumerUrl);
		authRequest.setIssuer(issuer);
		authRequest.setNameIDPolicy(nameIdPolicy);
		authRequest.setRequestedAuthnContext(requestedAuthnContext);
		authRequest.setID(randomId);
		authRequest.setVersion(SAMLVersion.VERSION_20);

		return authRequest;

	}

	private String encodeAuthnRequest(AuthnRequest authnRequest)
			throws MarshallingException, IOException {

		Marshaller marshaller = null;
		org.w3c.dom.Element authDOM = null;
		StringWriter requestWriter = null;
		String requestMessage = null;
		Deflater deflater = null;
		ByteArrayOutputStream byteArrayOutputStream = null;
		DeflaterOutputStream deflaterOutputStream = null;
		String encodedRequestMessage = null;

		marshaller = org.opensaml.Configuration.getMarshallerFactory()
				.getMarshaller(authnRequest); // object to DOM converter

		authDOM = marshaller.marshall(authnRequest); // converting to a DOM

		requestWriter = new StringWriter();
		XMLHelper.writeNode(authDOM, requestWriter);
		requestMessage = requestWriter.toString(); // DOM to string

		deflater = new Deflater(Deflater.DEFLATED, true);
		byteArrayOutputStream = new ByteArrayOutputStream();
		deflaterOutputStream = new DeflaterOutputStream(byteArrayOutputStream,
				deflater);
		deflaterOutputStream.write(requestMessage.getBytes()); // compressing
		deflaterOutputStream.close();

		encodedRequestMessage = Base64.encodeBytes(byteArrayOutputStream
				.toByteArray(), Base64.DONT_BREAK_LINES);
		encodedRequestMessage = URLEncoder.encode(encodedRequestMessage,
				"UTF-8").trim(); // encoding string

		return encodedRequestMessage;

	}

	public String processResponseMessage(String responseMessage) {

		XMLObject responseObject = null;

		try {

			responseObject = this.unmarshall(responseMessage);

		} catch (ConfigurationException e) {
			e.printStackTrace();
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
		} catch (SAXException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (UnmarshallingException e) {
			e.printStackTrace();
		}

		return this.getResult(responseObject);
	}

	private XMLObject unmarshall(String responseMessage)
			throws ConfigurationException, ParserConfigurationException,
			SAXException, IOException, UnmarshallingException {

		DocumentBuilderFactory documentBuilderFactory = null;
		DocumentBuilder docBuilder = null;
		Document document = null;
		Element element = null;
		UnmarshallerFactory unmarshallerFactory = null;
		Unmarshaller unmarshaller = null;

		DefaultBootstrap.bootstrap();

		documentBuilderFactory = DocumentBuilderFactory.newInstance();

		documentBuilderFactory.setNamespaceAware(true);

		docBuilder = documentBuilderFactory.newDocumentBuilder();

		document = docBuilder.parse(new ByteArrayInputStream(responseMessage
				.trim().getBytes())); // response to DOM

		element = document.getDocumentElement(); // the DOM element

		unmarshallerFactory = Configuration.getUnmarshallerFactory();

		unmarshaller = unmarshallerFactory.getUnmarshaller(element);

		return unmarshaller.unmarshall(element); // Response object

	}

	private String getResult(XMLObject responseObject) {

		Element ele = null;
		NodeList statusNodeList = null;
		Node statusNode = null;
		NamedNodeMap statusAttr = null;
		Node valueAtt = null;
		String statusValue = null;

		String[] word = null;
		String result = null;

		NodeList nameIDNodeList = null;
		Node nameIDNode = null;
		String nameID = null;

		// reading the Response Object
		ele = responseObject.getDOM();
		statusNodeList = ele.getElementsByTagName("samlp:StatusCode");
		statusNode = statusNodeList.item(0);
		statusAttr = statusNode.getAttributes();
		valueAtt = statusAttr.item(0);
		statusValue = valueAtt.getNodeValue();

		word = statusValue.split(":");
		result = word[word.length - 1];

		nameIDNodeList = ele.getElementsByTagNameNS(
				"urn:oasis:names:tc:SAML:2.0:assertion", "NameID");
		nameIDNode = nameIDNodeList.item(0);
		nameID = nameIDNode.getFirstChild().getNodeValue();

		result = nameID + ":" + result;

		return result;
	}

}
